
#ifndef KEYWORDSPOTTER_H_
#define KEYWORDSPOTTER_H_

#include <memory>
#include <string>
#include <vector>
#include <iomanip>
#include <unordered_set>

#include "MNN/Interpreter.hpp"

namespace tykws {

    struct Token {
        int timeStep;  //  token time step
        int id;        //  token id of vocab
        float prob;    //  token prob
    };

    struct CtcPrefixBeamSearchOptions {
        int blank = 0;                // blank id of vocab list.
        int first_beam_size = 3;
        int second_beam_size = 3;
    };

    struct PrefixHash {
        size_t operator()(const std::vector<int>& prefix) const {
            size_t hash_code = 0;
            // here we use KB&DR hash code
            for (int id : prefix) {
                hash_code = id + 31 * hash_code;
            }
            return hash_code;
        }
    };

    struct PrefixScore {             // for one prefix.
        float s = 0.0;               // blank ending score
        float ns = 0.0;              // none blank ending score
        std::vector<Token> nodes;
        float total_score() const { return ns + s; }
    };

    struct KeyWord {             // for keyword.
        float hit_score = 1.0;
        int start_frame = 0;
        int end_frame = 0;
        bool state = false; // is activated or not.
    };

    // Define decoding type.
    typedef enum {
        DECODE_GREEDY_SEARCH=0,
        DECODE_PREFIX_BEAM_SEARCH=1,
    }DECODE_TYPE;

    class KeywordSpotting {
    public:
        explicit KeywordSpotting(const std::string &model_path, DECODE_TYPE decode_type, int model_type);

        ~KeywordSpotting() { delete rTensor;}

        void Reset();

        void reset_value();

        void Forward(const std::vector<std::vector<float>> &feats,
                     std::vector<std::vector<float>> *prob);

        // function to load vocab from token.txt
        void readToken(const std::string& tokenFile) ;

        // set keyword
        void setKeyWord(const std::string& keyWord);

        void decode_keywords(std::vector<std::vector<float>>& probs, float hitScoreThr=0.0);

        // decoding alignments to predict sequence using greedy search.
        void decode_with_greedy_search(int offset,  std::vector<std::vector<float>>& probs);

        // decoding alignments to predict sequence using prefix beam search.
        void decode_ctc_prefix_beam_search(int offset, const std::vector<float> &prob);

        // find keyword
        void execute_detection(float hitScoreThr=0.1);

        // Token is keyword or not.
        bool isKeyword(int index);

        //update current hypotheses from proposed extensions.
        void UpdateHypotheses(const std::vector<std::pair<std::vector<int>, PrefixScore>>& hpys);

        // maxpooling keywords
        std::vector<std::string> mmaxpooling_keywords;

        //time stemp reset
        void stepClear();

        //keyword info;
        KeyWord kwsInfo;        

    private:
        // mnn runtime session
        std::shared_ptr<MNN::Interpreter> net;
        MNN::Session* session_net;

        /*inputs*/
        MNN::Tensor *xTensor;    
        MNN::Tensor *cacheTensor;    
        /*outputs*/
        MNN::Tensor *yTensor;
        MNN::Tensor *r_cacheTensor;
        /*cache*/
        MNN::Tensor *rTensor;

        // set mdoel type.
        int mmodel_type;

        //set decoder type.
        DECODE_TYPE mdecode_type;

        // vocab {token:index}
        std::unordered_map<std::string, int> mvocab;
        // keyword string
        std::string  mkey_word ;
        // keyword index set
        std::unordered_set<int> mkeyword_set;
        // keyword index list. handle same token in keyword.
        std::vector<int> mkeyword_token;

        // CTC alignments.
        std::vector<Token> alignments;
        // set of hypotheses with greed search.
        std::vector<Token> gd_cur_hyps;
        // set of hypotheses with prefix beam search.
        std::unordered_map<std::vector<int>, PrefixScore, PrefixHash> cur_hyps_;
        int total_frames=0;// frame offset, for absolute time

        //ctc prefix beam search
        const CtcPrefixBeamSearchOptions opts_={0, 3, 10};

        // silence time, 1s audio = 99frames melFbank. with default frame_shift(10ms) and frame_length(25ms).
        // Now we set silenceFrames = 3s * 99 = 297.
        //int mSilenceFrames = 297;

        //global Time step.
        int mGTimeStep = 0;

        bool activated = false;
    };

} // namespace tykws


#endif